Move external sampler interface to AbstractMCMC #2704

penelopeysm · 2025-10-24T13:24:11Z

Following on from:

Add AbstractMCMC.getstats and AbstractMCMC.requires_unconstrained_space AbstractMCMC.jl#182 Adding the interface function AbstractMCMC.getstats
Implement AbstractMCMC.getstats AdvancedHMC.jl#471 and implement AbstractMCMC.getstats AdvancedMH.jl#119 Implementing them for AdvancedHMC and AdvancedMH

this PR changes Turing's external sampler interface to exclusively use AbstractMCMC functions. With this PR, anyone who defines an external sampler will only need to depend on AbstractMCMC (which they presumably already do, because it is a sampler) and LogDensityProblems (which is already a dep of AbstractMCMC). No need for a Turing extension.

To be precise, it makes the following changes:

Previously where one had to define Turing.Inference.getparams, now one has to define AbstractMCMC.getparams.
Previously there was no way to include sampler stats in the resulting chain, now one can define AbstractMCMC.getstats.
The default for Turing.Inference.isgibbscomponent is changed to true, so that external sampler packages don't need to override it (unless absolutely necessary).

As an example implementation, this PR contains a test mock (note how it doesn't require a Turing dep):

Turing.jl/test/mcmc/external_sampler.jl

Lines 20 to 74 in 06752c4

    
           # Turing declares an interface for external samplers (see docstring for 
        
           # ExternalSampler). We should check that implementing this interface 
        
           # and only this interface allows us to use the sampler in Turing. 
        
           struct MyState{V<:AbstractVector} 
        
               params::V 
        
           end 
        
           AbstractMCMC.getparams(s::MyState) = s.params 
        
           AbstractMCMC.getstats(s::MyState) = (param_length=length(s.params),) 
        
           # externalsamplers must accept LogDensityModel inside their step function. 
        
           # By default Turing gives the externalsampler a LDF constructed with 
        
           # adtype=ForwardDiff, so we should expect that inside the sampler we can 
        
           # call both `logdensity` and `logdensity_and_gradient`. 
        
           # 
        
           # The behaviour of this sampler is to simply calculate logp and its 
        
           # gradient, and then return the same values. 
        
           # 
        
           # TODO: Do we also want to run ADTypeCheckContext to make sure that it is  
        
           # indeed using the adtype provided from Turing? 
        
           struct MySampler <: AbstractMCMC.AbstractSampler end 
        
           function AbstractMCMC.step( 
        
               rng::Random.AbstractRNG, 
        
               model::AbstractMCMC.LogDensityModel, 
        
               sampler::MySampler; 
        
               # This initial_params should be an AbstractVector because the model is just a 
        
               # LogDensityModel, not a DynamicPPL.Model 
        
               initial_params::AbstractVector, 
        
               kwargs..., 
        
           ) 
        
               # Step 1 
        
               ldf = model.logdensity 
        
               lp = LogDensityProblems.logdensity(ldf, initial_params) 
        
               @test lp isa Real 
        
               lp, grad = LogDensityProblems.logdensity_and_gradient(ldf, initial_params) 
        
               @test lp isa Real 
        
               @test grad isa AbstractVector{<:Real} 
        
               return nothing, MyState(initial_params) 
        
           end 
        
           function AbstractMCMC.step( 
        
               rng::Random.AbstractRNG, 
        
               model::AbstractMCMC.LogDensityModel, 
        
               sampler::MySampler, 
        
               state::MyState; 
        
               kwargs..., 
        
           ) 
        
               # Step >= 1 
        
               params = state.params 
        
               ldf = model.logdensity 
        
               lp = LogDensityProblems.logdensity(ldf, params) 
        
               @test lp isa Real 
        
               lp, grad = LogDensityProblems.logdensity_and_gradient(ldf, params) 
        
               @test lp isa Real 
        
               @test grad isa AbstractVector{<:Real} 
        
               return nothing, MyState(params) 
        
           end

github-actions · 2025-10-24T13:41:46Z

Turing.jl documentation for PR #2704 is available at:
https://TuringLang.github.io/Turing.jl/previews/PR2704/

codecov · 2025-10-24T13:47:33Z

Codecov Report

❌ Patch coverage is 88.23529% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.52%. Comparing base (be007f3) to head (d8dcd11).
⚠️ Report is 1 commits behind head on breaking.

Files with missing lines	Patch %	Lines
src/mcmc/gibbs.jl	71.42%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking    #2704      +/-   ##
============================================
+ Coverage     86.45%   86.52%   +0.06%     
============================================
  Files            21       21              
  Lines          1418     1410       -8     
============================================
- Hits           1226     1220       -6     
+ Misses          192      190       -2

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

penelopeysm · 2025-10-24T13:30:09Z

src/mcmc/external_sampler.jl

@@ -1,5 +1,5 @@
 """
-    ExternalSampler{S<:AbstractSampler,AD<:ADTypes.AbstractADType,Unconstrained}
+    ExternalSampler{Unconstrained,S<:AbstractSampler,AD<:ADTypes.AbstractADType}


Moved this type parameter earlier so that we can dispatch on it more easily.

penelopeysm · 2025-11-03T10:58:21Z

src/mcmc/external_sampler.jl


-1. Directly implement the `AbstractMCMC.step` methods for `DynamicPPL.Model`. This is the
-   most powerful option and is what Turing.jl's in-house samplers do. Implementing this
-   means that you can directly call `sample(model, MySampler(), N)`.
+1. Directly implement the `AbstractMCMC.step` methods for `DynamicPPL.Model`. That is to
+   say, implement `AbstractMCMC.step(rng::Random.AbstractRNG, model::DynamicPPL.Model,
+   sampler::MySampler; kwargs...)` and related methods. This is the most powerful option and
+   is what Turing.jl's in-house samplers do. Implementing this means that you can directly
+   call `sample(model, MySampler(), N)`.

-2. Implement a generic `AbstractMCMC.step` method for `AbstractMCMC.LogDensityModel`. This
-   struct wraps an object that obeys the LogDensityProblems.jl interface, so your `step`
+2. Implement a generic `AbstractMCMC.step` method for `AbstractMCMC.LogDensityModel` (the
+   same signature as above except that `model::AbstractMCMC.LogDensityModel`). This struct
+   wraps an object that obeys the LogDensityProblems.jl interface, so your `step`
   implementation does not need to know anything about Turing.jl or DynamicPPL.jl. To use
   this with Turing.jl, you will need to wrap your sampler: `sample(model,
   externalsampler(MySampler()), N)`.


I'll write proper docs about this separately (in the docs repo).

sunxd3

I see no issue

Decouple external sampler interface from Turing

06752c4

penelopeysm changed the base branch from main to breaking October 24, 2025 13:24

github-actions bot assigned penelopeysm Oct 24, 2025

penelopeysm added 3 commits October 24, 2025 14:30

Delete a todo note

8fec487

Changelog

46354fb

Temp point to AbstractMCMC feature branch

c7ffada

penelopeysm marked this pull request as draft October 24, 2025 13:39

penelopeysm commented Oct 24, 2025

View reviewed changes

penelopeysm added 6 commits October 24, 2025 16:39

Fix tests

eacffdb

Don't remove the keyword argument

9b30bea

Fix import order

0ebf0e8

Upstream getstats definitions to AdvancedMH/AdvancedHMC

c464b0c

remove sources

e99c154

Bump test deps (probably pointless, but eh)

d8dcd11

penelopeysm marked this pull request as ready for review November 4, 2025 15:47

penelopeysm commented Nov 4, 2025

View reviewed changes

penelopeysm requested a review from sunxd3 November 4, 2025 15:47

sunxd3 approved these changes Nov 5, 2025

View reviewed changes

penelopeysm merged commit f8a91a1 into breaking Nov 5, 2025
28 checks passed

penelopeysm deleted the py/tochains branch November 5, 2025 15:02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Move external sampler interface to AbstractMCMC #2704

Move external sampler interface to AbstractMCMC #2704

Uh oh!

penelopeysm commented Oct 24, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Oct 24, 2025

Uh oh!

codecov bot commented Oct 24, 2025 •

edited

Loading

Uh oh!

penelopeysm Oct 24, 2025

Uh oh!

penelopeysm Nov 3, 2025 •

edited

Loading

Uh oh!

sunxd3 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	# Turing declares an interface for external samplers (see docstring for
	# ExternalSampler). We should check that implementing this interface
	# and only this interface allows us to use the sampler in Turing.
	struct MyState{V<:AbstractVector}
	params::V
	end
	AbstractMCMC.getparams(s::MyState) = s.params
	AbstractMCMC.getstats(s::MyState) = (param_length=length(s.params),)

	# externalsamplers must accept LogDensityModel inside their step function.
	# By default Turing gives the externalsampler a LDF constructed with
	# adtype=ForwardDiff, so we should expect that inside the sampler we can
	# call both `logdensity` and `logdensity_and_gradient`.
	#
	# The behaviour of this sampler is to simply calculate logp and its
	# gradient, and then return the same values.
	#
	# TODO: Do we also want to run ADTypeCheckContext to make sure that it is
	# indeed using the adtype provided from Turing?
	struct MySampler <: AbstractMCMC.AbstractSampler end
	function AbstractMCMC.step(
	rng::Random.AbstractRNG,
	model::AbstractMCMC.LogDensityModel,
	sampler::MySampler;
	# This initial_params should be an AbstractVector because the model is just a
	# LogDensityModel, not a DynamicPPL.Model
	initial_params::AbstractVector,
	kwargs...,
	)
	# Step 1
	ldf = model.logdensity
	lp = LogDensityProblems.logdensity(ldf, initial_params)
	@test lp isa Real
	lp, grad = LogDensityProblems.logdensity_and_gradient(ldf, initial_params)
	@test lp isa Real
	@test grad isa AbstractVector{<:Real}
	return nothing, MyState(initial_params)
	end
	function AbstractMCMC.step(
	rng::Random.AbstractRNG,
	model::AbstractMCMC.LogDensityModel,
	sampler::MySampler,
	state::MyState;
	kwargs...,
	)
	# Step >= 1
	params = state.params
	ldf = model.logdensity
	lp = LogDensityProblems.logdensity(ldf, params)
	@test lp isa Real
	lp, grad = LogDensityProblems.logdensity_and_gradient(ldf, params)
	@test lp isa Real
	@test grad isa AbstractVector{<:Real}
	return nothing, MyState(params)
	end

Move external sampler interface to AbstractMCMC #2704

Move external sampler interface to AbstractMCMC #2704

Uh oh!

Conversation

penelopeysm commented Oct 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Oct 24, 2025

Uh oh!

codecov bot commented Oct 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

penelopeysm Oct 24, 2025

Choose a reason for hiding this comment

Uh oh!

penelopeysm Nov 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sunxd3 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

penelopeysm commented Oct 24, 2025 •

edited

Loading

codecov bot commented Oct 24, 2025 •

edited

Loading

penelopeysm Nov 3, 2025 •

edited

Loading